from typing import Dict, Tuple, List

import cv2
import numpy as np

from taskinfo.image_utils import ImageTransform
from taskinfo.table_process import bbox_decode, decode_by_ind, nms, bbox_post_process, sigmoid


class UvPreProcessor:
    """UV文档预处理器"""

    def __init__(self):
        self.inp_height = 768
        self.inp_width = 768
        self.mean = np.array([0.408, 0.447, 0.470], dtype=np.float32).reshape(1, 1, 3)
        self.std = np.array([0.289, 0.274, 0.278], dtype=np.float32).reshape(1, 1, 3)

    def __call__(self, img: np.ndarray) -> Tuple[np.ndarray, Dict]:
        """预处理图像"""
        if img is None or img.size == 0:
            raise ValueError("Input image is empty")

        height, width = img.shape[0:2]
        c = np.array([width / 2., height / 2.], dtype=np.float32)
        s = max(height, width) * 1.0

        trans_input = ImageTransform.get_affine_transform(
            c, s, 0, [self.inp_width, self.inp_height])
        inp_image = cv2.warpAffine(
            img,
            trans_input,
            (self.inp_width, self.inp_height),
            flags=cv2.INTER_LINEAR
        )

        # 添加数值范围检查
        inp_image = np.clip(inp_image, 0, 255)
        inp_image = ((inp_image / 255. - self.mean) / self.std).astype(np.float32)

        # 合并转置和reshape操作
        image = np.ascontiguousarray(inp_image.transpose(2, 0, 1).reshape(1, 3, self.inp_height, self.inp_width))
        meta = {
            'c': c,
            's': s,
            'input_height': self.inp_height,
            'input_width': self.inp_width,
            'out_height': self.inp_height // 4,
            'out_width': self.inp_width // 4
        }
        return image, meta


class UvPostProcessor:
    """UV文档后处理器"""
    SCORE_THRESHOLD = 0.3
    NMS_THRESHOLD = 0.3

    def __init__(self):
        self.rotate_map = {1: 2, 2: 1, 3: 0}  # 对应 cv2.rotate 的参数

    def __call__(self, model_outputs: Tuple[np.ndarray, ...], meta: Dict,
                 img: np.ndarray) -> Dict[str, np.ndarray]:
        """后处理模型输出"""
        hm, cls, ftype, wh, reg = model_outputs

        # 验证模型输出形状
        for output in model_outputs:
            if output is None or output.size == 0:
                raise ValueError("Model output is empty")

        angle_cls = sigmoid(cls)
        ftype_cls = sigmoid(ftype)
        hm = sigmoid(hm)
        bbox, inds = bbox_decode(hm, wh, reg=reg, K=10)
        angle_cls = decode_by_ind(angle_cls, inds, K=10)
        ftype_cls = decode_by_ind(ftype_cls, inds, K=10).astype(np.float32)

        bbox[0, :, 9] = angle_cls[0]
        bbox = np.concatenate((bbox, np.expand_dims(ftype_cls, axis=-1)), axis=-1)
        bbox = nms(bbox, self.NMS_THRESHOLD)

        if bbox.size == 0:
            return {
                "polygons": [],
                "scores": [],
                "output_imgs": [],
                "labels": [],
                "layout": np.array([])
            }

        bbox = bbox_post_process(bbox, [meta['c']], [meta['s']],
                                 meta['out_height'], meta['out_width'])

        res: List[np.ndarray] = []
        angle: List[int] = []
        sub_imgs: List[np.ndarray] = []
        ftype: List[int] = []
        score: List[float] = []

        for box in bbox[0]:
            if box[8] > self.SCORE_THRESHOLD:
                angle.append(int(box[9]))
                res.append(box[0:8])
                try:
                    sub_img = ImageTransform.crop_quad(img, res[-1].reshape(4, 2))
                    if angle[-1] in self.rotate_map:
                        sub_img = cv2.rotate(sub_img, self.rotate_map[angle[-1]])
                    sub_imgs.append(sub_img)
                except Exception:
                    continue
                ftype.append(int(box[10]))
                score.append(box[8])

        return {
            "polygons": res,
            "scores": score,
            "output_imgs": sub_imgs,
            "labels": angle,
            "layout": np.array(ftype) if ftype else np.array([])
        }
